package rbtree;

public class RBTree {
    static class RBTreeNode{
        public RBTreeNode parent;
        public RBTreeNode left;
        public RBTreeNode right;
        public int val;
        public COLOR color;
        public RBTreeNode(int val){
            this.val = val;
            this.color = COLOR.RED;
        }
    }
    //根节点
    public RBTreeNode root;
    //插入
    public boolean insert(int val){
        RBTreeNode node = new RBTreeNode(val);
        if(root == null){
            root = node;
            node.color = COLOR.BLACK;
            return true;
        }
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while(cur != null){
            if(val < cur.val){
                parent = cur;
                cur = cur.left;
            }else if(val == cur.val){
                return false;
            }else{
                parent = cur;
                cur = cur.right;
            }
        }
        if(val < parent.val){
            parent.left = node;
            node.parent = parent;
        }else{
            parent.right = node;
            node.parent = parent;
        }
        cur = node;
        while(parent != null && parent.color == COLOR.RED){
            RBTreeNode grandParent = parent.parent;
            if(parent == grandParent.left){
                RBTreeNode uncle = grandParent.right;
                if(uncle != null && uncle.color == COLOR.RED){
                    //uncle为红色
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    if(grandParent != root){
                        grandParent.color = COLOR.RED;
                    }
                    cur = grandParent;
                    parent = cur.parent;
                }else{
                    //uncle不存在或为黑色
                    if(cur == parent.right){
                        rotateLeft(parent);
                        RBTreeNode mid = parent;
                        parent = cur;
                        cur = mid;
                    }
                    rotateRight(grandParent);
                    parent.color = COLOR.BLACK;
                    grandParent.color = COLOR.RED;
                }
            }else{
                RBTreeNode uncle = grandParent.left;
                if(uncle !=null && uncle.color == COLOR.RED){
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    if(grandParent != root){
                        grandParent.color = COLOR.RED;
                    }
                    cur = grandParent;
                    parent = grandParent.parent;
                }else{
                    if(parent.left == cur){
                        rotateRight(parent);
                        RBTreeNode mid = parent;
                        parent = cur;
                        cur = mid;
                    }
                    rotateLeft(grandParent);
                    parent.color = COLOR.BLACK;
                    grandParent.color = COLOR.RED;
                }
            }
        }
        return true;
    }
    private void rotateLeft(RBTreeNode parent){
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;
        parent.right = subRL;
        if(subRL != null){
            subRL.parent = parent;
        }
        subR.left = parent;
        RBTreeNode pParent = parent.parent;
        parent.parent = subR;
        if(root == parent){
            root = subR;
        }else{
            if(pParent.left ==parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
        }
        subR.parent = pParent;
    }
    private void rotateRight(RBTreeNode parent){
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;
        parent.left = subLR;
        if(subLR != null){
            subLR.parent = parent;
        }
        subL.right = parent;
        RBTreeNode pParent = parent.parent;
        parent.parent = subL;
        if(root == parent){
            root = subL;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
        }
        subL.parent = pParent;
    }
    //判断一棵树是否是红黑树
    public boolean isRBTree(RBTreeNode root){
        if(root == null){
            return true;
        }
        if(root.color != COLOR.BLACK){
            System.out.println("根节点不是黑色，所以不是红黑树");
            return false;
        }
        if(checkRedColor(root) == false){
            System.out.println("存在两个连续的红色节点，所以不是红黑树");
            return false;
        }
        RBTreeNode node = root;
        int blackNum = 0;
        while(node != null){
            if(node.color == COLOR.BLACK){
                blackNum++;
            }
            node = node.left;
        }
        if(checkBlackNum(root,0,blackNum) == false){
            System.out.println("存在两条路径上黑色节点个数不同的情况，所以不是红黑树");
            return false;
        }
        return true;
    }
    private boolean checkRedColor(RBTreeNode root){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.RED){
            RBTreeNode parent = root.parent;
            if(parent.color == COLOR.RED){
                return false;
            }
        }
        return checkRedColor(root.left)&&checkRedColor(root.right);
    }
    private boolean checkBlackNum(RBTreeNode root,int pathNum,int blackNum){
        if(root == null){
            return pathNum == blackNum;
        }
        if(root.color == COLOR.BLACK){
            pathNum++;
        }
        return checkBlackNum(root.left,pathNum,blackNum) && checkBlackNum(root.right,pathNum,blackNum);
    }
    public void inOrder(RBTreeNode root){
        if(root == null){
            return;
        }
        inOrder(root.left);
        System.out.println(root.val+" ");
        inOrder(root.right);
    }
}
